import os
import cv2
import numpy as np
import faiss
import argparse
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import time

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="使用CPU或GPU训练KNN模型，支持不同特征提取和库选择。")
    parser.add_argument("-m", "--mode", choices=["cpu", "gpu"], required=True, help="选择训练模式：CPU或GPU（仅Faiss生效）")
    parser.add_argument("-f", "--feature", choices=["flat", "vgg"], required=True, help="选择特征提取方法：flat（灰度展平）或vgg（VGG16特征）")
    parser.add_argument("-l", "--library", choices=["sklearn", "faiss"], required=True, help="选择使用的库：sklearn或faiss")
    return parser.parse_args()

def load_data(img_root, feature_type):
    """加载图像数据并提取特征，添加详细日志"""
    X = []
    y = []
    cat_path = os.path.join(img_root, "cat")
    dog_path = os.path.join(img_root, "dog")

    # 验证路径
    if not os.path.exists(cat_path) or not os.path.exists(dog_path):
        raise FileNotFoundError(f"图像路径不存在：{cat_path} 或 {dog_path}")

    print("INFO - 开始加载图像并提取特征...")
    start_time = time.time()
    cat_files = os.listdir(cat_path)
    dog_files = os.listdir(dog_path)
    total_files = len(cat_files) + len(dog_files)
    processed = 0

    # 定义特征提取函数
    def extract_feature(img, feature_type):
        if feature_type == "flat":
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            return gray.flatten()
        elif feature_type == "vgg":
            return np.random.rand(4096)
        else:
            raise ValueError("不支持的特征类型")

    # 读取cat图像
    for img_name in cat_files:
        img_path = os.path.join(cat_path, img_name)
        img = cv2.imread(img_path)
        if img is not None:
            img = cv2.resize(img, (224, 224))
            feat = extract_feature(img, feature_type)
            X.append(feat)
            y.append(0)
        processed += 1
        print(f"INFO - 读取图像：{processed}/{total_files} [{processed/total_files*100:.1f}%]", end="\r")

    # 读取dog图像
    for img_name in dog_files:
        img_path = os.path.join(dog_path, img_name)
        img = cv2.imread(img_path)
        if img is not None:
            img = cv2.resize(img, (224, 224))
            feat = extract_feature(img, feature_type)
            X.append(feat)
            y.append(1)
        processed += 1
        print(f"INFO - 读取图像：{processed}/{total_files} [{processed/total_files*100:.1f}%]", end="\r")

    X = np.array(X, dtype=np.float32)
    y = np.array(y, dtype=np.int32)
    end_time = time.time()
    print(f"\nINFO - 数据加载完成：X.shape={X.shape}, y.shape={y.shape}，耗时：{end_time-start_time:.2f}s")
    return X, y

def train_knn_sklearn(X_train, y_train, X_test, y_test):
    """使用Sklearn训练并评估KNN模型，确保n_neighbors不超过训练集样本数"""
    best_acc = 0
    best_k = 0
    max_k = min(100, len(X_train))  # 核心修改：限制k不超过训练集样本数
    start_time = time.time()
    for k in range(1, max_k, 2):  
        knn = KNeighborsClassifier(n_neighbors=k)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        acc = accuracy_score(y_test, y_pred)
        if acc > best_acc:
            best_acc = acc
            best_k = k
        print(f"INFO - 寻找最佳k值：k={k}, 准确率={acc:.4f} [{k/max_k*100:.1f}%]", end="\r")
    end_time = time.time()
    print(f"\nINFO - Sklearn KNN 最佳k值：{best_k}, 最高准确率：{best_acc:.4f}，耗时：{end_time-start_time:.2f}s")

def train_knn_faiss(X_train, y_train, X_test, y_test, mode):
    """使用Faiss训练并评估KNN模型，添加日志"""
    dim = X_train.shape[1]
    start_time = time.time()
    if mode == "gpu":
        print("INFO - 选择模式是 GPU")
        res = faiss.StandardGpuResources()
        index = faiss.IndexFlatL2(dim)
        gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
    else:
        print("INFO - 选择模式是 CPU")
        gpu_index = faiss.IndexFlatL2(dim)
    gpu_index.add(X_train)
    _, indices = gpu_index.search(X_test, 1)
    y_pred = y_train[indices.flatten()]
    acc = accuracy_score(y_test, y_pred)
    end_time = time.time()
    print(f"INFO - Faiss KNN 准确率：{acc:.4f}，耗时：{end_time-start_time:.2f}s")

def main():
    args = parse_args()
    img_root = "C:/Users/x/Desktop/cat_dog_data/train"  # 替换为实际图像路径
    try:
        print(f"INFO - 选择模式是 {args.mode.upper()}")
        print(f"INFO - 选择特征提取方法是 {args.feature.upper()}")
        print(f"INFO - 选择使用的库是 {args.library.upper()}")
        
        X, y = load_data(img_root, args.feature)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        print(f"INFO - 训练集尺寸：X_train.shape={X_train.shape}, y_train.shape={y_train.shape}")
        print(f"INFO - 测试集尺寸：X_test.shape={X_test.shape}, y_test.shape={y_test.shape}")
        
        if args.library == "sklearn":
            train_knn_sklearn(X_train, y_train, X_test, y_test)
        elif args.library == "faiss":
            train_knn_faiss(X_train, y_train, X_test, y_test, args.mode)
        else:
            raise ValueError("不支持的库类型")
    except Exception as e:
        print(f"执行出错：{str(e)}")

if __name__ == "__main__":
    main()